import onnx
from . import utils

def analyze(onnxfile, **kwargs):
    model = onnx.load(onnxfile)

    dict_pad = {}
    dict_pool = {}
    dict_mul = {}

    for node_id, node in enumerate(model.graph.node):
        if node.op_type == 'Pad':
            dict_pad['input'] = node.input
            dict_pad['output'] = node.output
            dict_pad['name'] = node.name

        if node.op_type == 'MaxPool':
            if len(dict_pad) > 0 and node.input == dict_pad['output']:
                dict_pool['input'] = node.input
                dict_pool['output'] = node.output
                dict_pool['id'] = node_id
                print('Warning: got pad+pool pair, pad:', dict_pad['input'], dict_pad['output'])
                print('Warning: got pad+pool pair, pool:', dict_pool['input'], dict_pool['output'])

                return False
            else:
                dict_pad = {}    

    origin_model = kwargs['origin_model']
    model = onnx.load(origin_model)

    dict_pad = {}
    dict_pool = {}
    dict_mul = {}

    got_pad_pool = False

    pad_pool_list = []

    for node_id, node in enumerate(model.graph.node):
        #print(node_id, ", name:", node.name, ", input:", node.input, ", output:", node.output,  \
        #         ", op:", node.op_type, ', len(input):', len(node.input))

        if node.op_type == 'Pad':
            dict_pad['input'] = node.input
            dict_pad['output'] = node.output
            dict_pad['name'] = node.name

        if node.op_type == 'MaxPool' or node.op_type == 'AveragePool':
            if len(dict_pad) > 0 and node.input == dict_pad['output']:
                dict_pool['input'] = node.input
                dict_pool['output'] = node.output
                dict_pool['id'] = node_id
                print('got pad+pool pair, pad:', dict_pad['input'], dict_pad['output'])
                print('got pad+pool pair, pool:', dict_pool['input'], dict_pool['output'])
                pads = []

                got_pad_pool = True

                for init in model.graph.initializer:
                    if init.name == dict_pad['input'][1]:
                        print('got init(pads):', init.name)
                        dtype = init.data_type
                        np_dtype = utils.convert_ort_type_2_np(dtype)
                        if init.raw_data:
                            params_list = np.fromstring(init.raw_data, dtype=np_dtype)
                            for p in params_list:
                                pads.append(p)
                        else:
                            data_list = utils.get_data_list(dtype, init)
                            for p in data_list:
                                pads.append(p)
                    elif init.name == dict_pad['input'][2]:
                        print('got init(constane_value):', init.name)  

                pads_real = [pads[2], pads[3], pads[6], pads[7]]

                d = {}
                d['pool_name'] = node.name
                d['pads'] = pads_real
                d['pad_name'] = dict_pad['name']

                pad_pool_list.append(d)

                dict_pad = {}
                dict_pool = {}
            else:
                #print('clear pad dict')
                dict_pad = {}    

    if got_pad_pool == True:
        model = onnx.load(onnxfile)

        dict_pad = {}
        dict_pool = {}
        dict_mul = {}

        got_pad_pool_ = False

        pad_pool_list_ = []

        for node_id, node in enumerate(model.graph.node):
            if node.op_type == 'MaxPool' or node.op_type == 'AveragePool':
                for e in pad_pool_list:
                    if node.name == e['pool_name']:
                        for attr in node.attribute:
                            if attr.name == 'pads':
                                if attr.ints != e['pads']:
                                    print('pool attr pads:', attr.ints, ', Pad attr pads:', e['pads'])
                                    return False
  
    return True

